#ifndef UnstructMG_MG_HPP
#define UnstructMG_MG_HPP

#include "MG_types.hpp"
#include "operator.hpp"

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
class UnstructMG : public Solver<idx_t, setup_t, calc_t, dof> {
public:
    UnstructMG(const std::string _config_path);
    ~UnstructMG();

    // 实际solve时的层数，刨除了那些未被启用的粗化层
    idx_t num_levs = 0;
    // solve阶段的cycle类型：1 for V, 2 for W
    idx_t cycle_type = 1;
    // 在V-cycle中前平滑和后平滑的次数
    idx_t num_grid_sweeps[2] = {1, 1};

    idx_t * coarsen_types = nullptr;
    idx_t * coarsen_passes = nullptr;
    setup_t * coarsen_thetas = nullptr;
    setup_t * coarsen_smwgts = nullptr;// 光滑聚合得权重
    idx_t * coarsen_rstrs = nullptr;/// 如果用光滑聚合计算粗矩阵，限制矩阵用P的转置(1)还是Pbar的转置(0)

    // 每一层的smoother的类型
    std::string * relax_types = nullptr;
    // 每一层的smoother的权重系数
    calc_t * relax_wgts = nullptr;
    // 在solve时得插值和限制使用plain(0)还是光滑后的(1)
    idx_t * apply_RP = nullptr;

    const par_CSRMatrix<idx_t, setup_t, setup_t, dof> * base_A_ptr = nullptr;
    void SetOperator(const Operator<idx_t, setup_t, dof> & op);


    /// @brief 处理配置参数的文件
    /// @param _config_file 文件名
    void ParseConfig(const std::string _config_file);

    void SetCycleType(const idx_t _cycle_type);

protected:
    void Mult_ipl(const par_Vector<idx_t, calc_t, dof> & b, par_Vector<idx_t, calc_t, dof> & x) const;

    /// @brief 建立第ilev层的smoother
    /// @param ilev 层号
    /// @param A 第ilev层的矩阵
    void SetupSmoother(const idx_t ilev, const par_CSRMatrix<idx_t, setup_t, setup_t, dof> & A);

    /// @brief 执行第ilev层的前平滑
    /// @param ilev 层号
    void preSmooth(const idx_t ilev) const;
    /// @brief 执行第ilev层的后平滑
    /// @param ilev 层号
    void postSmooth(const idx_t ilev) const;
    /// @brief 在第ilev层内计算残差
    /// @param ilev 层号
    void calcResi(const idx_t ilev) const;
    /// @brief 将第ilev层的残差向第ilev+1层进行限制，并作为下一层求解的右端项
    /// @param ilev 层号
    void rstrResi(const idx_t ilev) const;
    /// @brief 将第ilev+1层的向量值插值回到第ilev层，以此更新第ilev层的向量
    /// @param ilev 层号
    void prlgErr(const idx_t ilev) const;
    /// @brief Solve阶段执行V-cycle过程
    void V_Cycle() const ;
    /// @brief Solve阶段执行W-cycle过程
    void W_Cycle() const ;

    void truncate();
// private:
public:
    std::vector<idx_t> skipped;
    // 执行粗化时生成的粗化层数
    idx_t coarse_nlevs = 0;
    // 多重网格是否拥有最细层的内存所有权
    bool own_A0 = false;

    // 每一层网格的向量
    par_Vector<idx_t, calc_t, dof> ** U_array = nullptr;
    par_Vector<idx_t, calc_t, dof> ** F_array = nullptr;
    par_Vector<idx_t, calc_t, dof> ** aux_arr = nullptr;
    // 每一层网格的矩阵
    par_CSRMatrix<idx_t, calc_t, calc_t, dof> ** A_array_high = nullptr;
    par_CSRMatrix<idx_t, data_t, calc_t, dof> ** A_array_low  = nullptr;
    // 每一层的smoother
    Solver<idx_t, setup_t, calc_t, dof> ** smoother = nullptr;
    // 每相邻两层间的限制算子 和插值算子
    par_CSRMatrix<idx_t, data_t, calc_t, 1> ** restrictor = nullptr;
    par_CSRMatrix<idx_t, data_t, calc_t, 1> ** prolngator = nullptr;
    par_CSRMatrix<idx_t, data_t, calc_t, dof> ** rstr_smth = nullptr;
    par_CSRMatrix<idx_t, data_t, calc_t, dof> ** prlg_smth = nullptr;

};

typedef UnstructMG<int, double, double, double, 1> UnstructMG_I32All64_DOF1;
typedef UnstructMG<int, double, double, double, 2> UnstructMG_I32All64_DOF2;
typedef UnstructMG<int, double, double, double, 3> UnstructMG_I32All64_DOF3;
typedef UnstructMG<int, double, double, double, 4> UnstructMG_I32All64_DOF4;

typedef UnstructMG<int, float , double, double, 1> UnstructMG_I32K64P64D32_DOF1;
typedef UnstructMG<int, float , double, double, 2> UnstructMG_I32K64P64D32_DOF2;
typedef UnstructMG<int, float , double, double, 3> UnstructMG_I32K64P64D32_DOF3;
typedef UnstructMG<int, float , double, double, 4> UnstructMG_I32K64P64D32_DOF4;

typedef UnstructMG<int, float , double, float , 1> UnstructMG_I32K64P32D32_DOF1;
typedef UnstructMG<int, float , double, float , 2> UnstructMG_I32K64P32D32_DOF2;
typedef UnstructMG<int, float , double, float , 3> UnstructMG_I32K64P32D32_DOF3;
typedef UnstructMG<int, float , double, float , 4> UnstructMG_I32K64P32D32_DOF4;

typedef UnstructMG<int, float , float , float , 1> UnstructMG_I32All32_DOF1;
typedef UnstructMG<int, float , float , float , 2> UnstructMG_I32All32_DOF2;
typedef UnstructMG<int, float , float , float , 3> UnstructMG_I32All32_DOF3;
typedef UnstructMG<int, float , float , float , 4> UnstructMG_I32All32_DOF4;
#ifdef USE_FP16
typedef UnstructMG<int, __fp16, double, float , 1> UnstructMG_I32K64P32D16_DOF1;
typedef UnstructMG<int, __fp16, double, float , 2> UnstructMG_I32K64P32D16_DOF2;
typedef UnstructMG<int, __fp16, double, float , 3> UnstructMG_I32K64P32D16_DOF3;
typedef UnstructMG<int, __fp16, double, float , 4> UnstructMG_I32K64P32D16_DOF4;

typedef UnstructMG<int, __fp16, float , float , 1> UnstructMG_I32K32P32D16_DOF1;
typedef UnstructMG<int, __fp16, float , float , 2> UnstructMG_I32K32P32D16_DOF2;
typedef UnstructMG<int, __fp16, float , float , 3> UnstructMG_I32K32P32D16_DOF3;
typedef UnstructMG<int, __fp16, float , float , 4> UnstructMG_I32K32P32D16_DOF4;
#endif

template<typename idx_t, typename data_t, typename setup_t1, typename setup_t2,
    typename calc_t1, typename calc_t2, int dof>
void BorrowMG(const UnstructMG<idx_t, data_t, setup_t1, calc_t1, dof> * src,
                    UnstructMG<idx_t, data_t, setup_t2, calc_t2, dof> * & dst);
extern template void BorrowMG(const UnstructMG_I32K64P64D32_DOF1 *, UnstructMG_I32K64P32D32_DOF1 * &);
extern template void BorrowMG(const UnstructMG_I32K64P64D32_DOF2 *, UnstructMG_I32K64P32D32_DOF2 * &);
extern template void BorrowMG(const UnstructMG_I32K64P64D32_DOF3 *, UnstructMG_I32K64P32D32_DOF3 * &);
extern template void BorrowMG(const UnstructMG_I32K64P64D32_DOF4 *, UnstructMG_I32K64P32D32_DOF4 * &);

extern template void BorrowMG(const UnstructMG_I32K64P64D32_DOF1 *, UnstructMG_I32All32_DOF1 * &);
extern template void BorrowMG(const UnstructMG_I32K64P64D32_DOF2 *, UnstructMG_I32All32_DOF2 * &);
extern template void BorrowMG(const UnstructMG_I32K64P64D32_DOF3 *, UnstructMG_I32All32_DOF3 * &);
extern template void BorrowMG(const UnstructMG_I32K64P64D32_DOF4 *, UnstructMG_I32All32_DOF4 * &);

template<typename idx_t, typename data_t, typename setup_t, typename calc_t2, int dof>
void ReturnMG(UnstructMG<idx_t, data_t, setup_t, calc_t2, dof> * & dst);
extern template void ReturnMG(UnstructMG_I32K64P32D32_DOF1 * &);
extern template void ReturnMG(UnstructMG_I32K64P32D32_DOF2 * &);
extern template void ReturnMG(UnstructMG_I32K64P32D32_DOF3 * &);
extern template void ReturnMG(UnstructMG_I32K64P32D32_DOF4 * &);


#endif
